# 导入llama_index核心模块中的StorageContext、load_index_from_storage、VectorStoreIndex和Document类
from llama_index.core import (
    StorageContext,
    load_index_from_storage,
    VectorStoreIndex,
    Document,
)
import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore

# 初始化chromadb客户端，数据将保存到指定的目录
db = chromadb.PersistentClient(path="./chroma_db")
# 获取或创建llamaindex集合
llamaindex_collection = db.get_or_create_collection("llamaindex")
# 创建ChromaVectorStore的实例，实现llamaindex和chromadb对接
vector_store = ChromaVectorStore(chroma_collection=llamaindex_collection)
# 创建chromadb的存储上下文，管理索引的存储配置 如向量存储
storage_context = StorageContext.from_defaults(vector_store=vector_store)
# 直接从向量数据库加载索引，无需重新处理文档
index = VectorStoreIndex.from_vector_store(
    vector_store, storage_context=storage_context
)
query_engine = index.as_query_engine()
response = query_engine.query("LlamaIndex的主要功能是什么？")
print(response)
